from dataclasses import dataclass
from typing import List, Dict, Optional, ClassVar, Literal


@dataclass(frozen=True)
class RobotFeaturesBase:
    # 子类需要重写
    motors: ClassVar[List[str]] = []
    cameras: ClassVar[List[str]] = []

    # 缓存生成的 features
    _features_cache: ClassVar[Optional[Dict]] = None

    fps: ClassVar[int] = 30

    @classmethod
    def features(
        cls,
        has_velocity: bool = False,
        has_effort: bool = False,
        image_dtype: Literal["video", "image"] = "video",
    ) -> Dict:
        """
        返回机器人特征字典
        """
        if cls._features_cache is not None:
            return cls._features_cache

        features = {
            "observation.state": {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            },
            "action": {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            },
        }

        if has_velocity:
            features["observation.velocity"] = {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            }

        if has_effort:
            features["observation.effort"] = {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            }

        for cam in cls.cameras:
            features[f"observation.images.{cam}"] = {
                "dtype": image_dtype,
                "shape": (3, 480, 640),
                "names": ["channels", "height", "width"],
            }

        cls._features_cache = features
        return features


@dataclass(frozen=True)
class SO100Features(RobotFeaturesBase):
    motors = [
        "waist",
        "shoulder",
        "elbow",
        "forearm_roll",
        "wrist_angle",
        "gripper",
    ]
    cameras = ["cam_wrist", "cam_top"]


@dataclass(frozen=True)
class SO100DualFeatures(RobotFeaturesBase):
    motors = [
        "left_waist",
        "left_shoulder",
        "left_elbow",
        "left_forearm_roll",
        "left_wrist_angle",
        "left_gripper",
        "right_waist",
        "right_shoulder",
        "right_elbow",
        "right_forearm_roll",
        "right_wrist_angle",
        "right_gripper",
    ]
    cameras = ["cam_left", "cam_head", "cam_right"]


@dataclass(frozen=True)
class AirbotDualFeatures(RobotFeaturesBase):
    motors = [
        "right_waist",
        "right_shoulder",
        "right_elbow",
        "right_forearm_roll",
        "right_wrist_angle",
        "right_wrist_rotate",
        "right_gripper",
        "left_waist",
        "left_shoulder",
        "left_elbow",
        "left_forearm_roll",
        "left_wrist_angle",
        "left_wrist_rotate",
        "left_gripper",
    ]
    eef = [
        "right_x",
        "right_y",
        "right_z",
        "right_roll",
        "right_pitch",
        "right_yaw",
        "right_gripper",
        "left_x",
        "left_y",
        "left_z",
        "left_roll",
        "left_pitch",
        "left_yaw",
        "left_gripper",
    ]
    cameras = ["cam_left", "cam_head", "cam_right"]

    fps = 20

    @classmethod
    def features(
        cls,
        has_velocity: bool = False,
        has_effort: bool = False,
        image_dtype: Literal["video", "image"] = "video",
    ) -> Dict:
        """
        返回机器人特征字典
        """
        if cls._features_cache is not None:
            return cls._features_cache

        features = {
            "observation.state": {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            },
            "observation.state_gpos": {
                "dtype": "float32",
                "shape": (len(cls.eef),),
                "names": cls.eef,
            },
            "action": {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            },
            "action_gpos": {
                "dtype": "float32",
                "shape": (len(cls.eef),),
                "names": cls.eef,
            },
        }

        if has_velocity:
            features["observation.velocity"] = {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            }

        if has_effort:
            features["observation.effort"] = {
                "dtype": "float32",
                "shape": (len(cls.motors),),
                "names": cls.motors,
            }

        for cam in cls.cameras:
            features[f"observation.images.{cam}"] = {
                "dtype": image_dtype,
                "shape": (3, 480, 640),
                "names": ["channels", "height", "width"],
            }

        cls._features_cache = features
        return features


ROBOT_MAP = {"so100": SO100Features, "airbot_dual": AirbotDualFeatures}
